from scipy.io import loadmat
import numpy as np
from keras.models import Sequential
from keras.layers import LSTM, Dense
from keras.callbacks import ModelCheckpoint, EarlyStopping, LearningRateScheduler

data = loadmat('S10_E1_A3.mat')
X = data.get('emg')
Y = data.get('glove')
X = np.expand_dims(X, 1) * 1e3
Y = Y

def scheduler(epoch, lr):
    if epoch < 50:
        return 1e-3
    elif epoch < 100:
        return 1e-4
    else:
        return 1e-5


model = Sequential()
model.add(LSTM(128, input_shape=(X.shape[1], X.shape[-1]), return_sequences=True))
model.add(LSTM(128))
model.add(Dense(Y.shape[-1]))
model.compile(loss='mae', optimizer='adam')
checkpoint = ModelCheckpoint('lstm.h5', monitor="val_loss", verbose=1, save_best_only=True)
earlystop = EarlyStopping(monitor="val_loss", patience=100, verbose=1)
learn_rate = LearningRateScheduler(schedule=scheduler, verbose=1)

# train LSTM
model.fit(X, Y, epochs=1000, batch_size=128, validation_split=0.2, verbose=1,
          callbacks=[checkpoint, earlystop, learn_rate])


